package com.any0ne.loadDemo;

import com.sun.tools.attach.AgentInitializationException;
import com.sun.tools.attach.AgentLoadException;
import com.sun.tools.attach.AttachNotSupportedException;
import com.sun.tools.attach.VirtualMachine;
import sun.jvmstat.monitor.*;
import java.io.IOException;
import java.lang.management.ManagementFactory;
import java.lang.management.RuntimeMXBean;
import java.net.URISyntaxException;
import java.util.HashSet;
import java.util.Set;

/**
 * Author: 4ut15m.4ny0neSec
 * Desc:
 * 这个类用来加载agent
 * */
public class Load {
    // simple-java-agent的根目录与默认的agent目录
    private static String projectRoot = Load.class.getResource("/").getPath().split("loadDemo")[0];
    private static String defaultAgentPath = "agentdemo/target/";
    private static String defaultAgent = "agentdemo-1.0-SNAPSHOT.jar";

    public static void sayHello() {
        System.out.println("Hello");
    }

    public static void main(String[] args) {
//        String defaultFile = Load.projectRoot + Load.defaultAgentPath + Load.defaultAgent;
//        if(System.getProperty("os.name").toLowerCase().startsWith("windows")){
//            defaultFile = defaultFile.substring(1, defaultFile.length());
//        }
//        System.out.println("This function is main.");
        sayHello();
        /**
         * 这个方法会打印当前所有的java进程及其pid.
         * */
//        allPid();
        /**
         * 需要注意的是，我这里给的例子是让程序自己以attach的方式加载agent。
         * 但实际上，我们也可以用一个程序把agent给attach到另一个程序上。
         * 举个例子：
         * 先在目前这个Load类的主方法中加入sleep 50s的代码
         * 再重新建一个idea项目，随便弄个类，把我给的这个attach方法移植过去。
         * 然后运行Load类的main方法
         * 再想办法获取这个Load类程序的pid
         * 再在新类中填入pid以及agent路径，运行，就可以看到Load类中被加载了agent输出了已加载的所有类
         * */
//        attach(null, defaultFile);
    }



    /**
     * 这个方法是用来以attach方法加载agent的,我们需要提前知道目标java进程的pid
     * agentPath是绝对路径,比如C:\Users\4ut15m\Desktop\Documents\Mine\agent\agentdemo\target\agentdemo-1.0-SNAPSHOT.jar
     */
    public static void attach(String pid, String agentPath) {
        String tmp;

        if (pid == null || pid.equals("")) {
            tmp = currentPid();
        } else {
            tmp = pid;
        }

        try {
            VirtualMachine virtualMachine = VirtualMachine.attach(tmp);
            virtualMachine.loadAgent(agentPath);
        } catch (AttachNotSupportedException e) {
            System.out.println("Attach error.");
            e.printStackTrace();
        } catch (IOException e) {
            e.printStackTrace();
        } catch (AgentLoadException e) {
            System.out.println("Agent load error");
            e.printStackTrace();
        } catch (AgentInitializationException e) {
            System.out.println("Agent initial error");
            e.printStackTrace();
        }
    }

    /**
     * 获取当前程序的pid
     */
    public static String currentPid() {
        RuntimeMXBean runtimeMXBean = ManagementFactory.getRuntimeMXBean();
        String pid = runtimeMXBean.getName().split("@")[0];
        return pid;
    }

    /**
     * 打印目前所有的java进行pid
     */
    public static void allPid() {
        // 获取监控主机
        MonitoredHost local = null;
        try {
            local = MonitoredHost.getMonitoredHost("localhost");
            // 取得所有在活动的虚拟机集合
            Set<?> vmlist = new HashSet<Object>(local.activeVms());
            // 遍历集合，输出PID和进程名
            for (Object process : vmlist) {
                MonitoredVm vm = local.getMonitoredVm(new VmIdentifier("//" + process));
                // 获取类名
                String processname = MonitoredVmUtil.mainClass(vm, true);
                System.out.println("进程名: "+processname + " Pid: " + process);
            }
        } catch (MonitorException e) {
            e.printStackTrace();
        } catch (URISyntaxException e) {
            e.printStackTrace();
        }

    }

}